package pers.whj.security;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;

public class RSA {
    public static final String ALGORITHM_NAME = "RSA";
    public static final String MD5_RSA = "MD5withRSA";

    public static enum KeyLength {
        LENGTH_256,
        LENGTH_512,
        LENGTH_1024,
        LENGTH_2048;
    }

    /**
     * encrypt the plain text
     *
     * @param plainText text to encode, should not be too long
     * @param publicKey public key
     * @param keyLength key length
     * @return the cipher text and encoded by Base64
     */
    public static String encrypt(String plainText, PublicKey publicKey, KeyLength keyLength) throws Exception {
        byte[] arr = plainText.getBytes(StandardCharsets.UTF_8);
        return Arrays.toString(encrypt(arr, 0, arr.length, publicKey, keyLength));
    }

    /**
     * encrypt the plain text
     *
     * @param plainText text to encode, should not be too long
     * @param publicKey public key
     * @param keyLength key length
     * @return the cipher text and encoded by Base64
     */
    public static byte[] encrypt(byte[] plainText, int offset, int length, PublicKey publicKey, KeyLength keyLength) throws Exception {
        checkKeyLength(keyLength);
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        int len = offset + length;
        int start = offset;
        byte[] buffer;
        int i = 0;
        final int MAX = getMaxEncryptLength(keyLength);
        while (len > start) {
            if (len - start > MAX) {
                buffer = cipher.doFinal(plainText, start, MAX);
            } else {
                buffer = cipher.doFinal(plainText, start, len - start);
            }
            outputStream.write(buffer, 0, buffer.length);
            i++;
            start = offset + i * MAX;
        }
        buffer = outputStream.toByteArray();
        outputStream.close();
        return Base64.getEncoder().encode(buffer);
    }

    /**
     * decrypt the cipher text
     *
     * @param cipherText text to decode, encoded by Base64
     * @param privateKey private key
     * @param keyLength  key length
     * @return the plain text
     */
    public static String decrypt(String cipherText, PrivateKey privateKey, KeyLength keyLength) throws Exception {
        byte[] arr = cipherText.getBytes(StandardCharsets.UTF_8);
        return new String(decrypt(arr, 0, arr.length, privateKey, keyLength), StandardCharsets.UTF_8);
    }

    /**
     * decrypt the cipher text
     *
     * @param cipherText text to decode, encoded by Base64
     * @param privateKey private key
     * @param keyLength  key length
     * @return the plain text
     */
    public static byte[] decrypt(byte[] cipherText, int offset, int length, PrivateKey privateKey, KeyLength keyLength) throws Exception {
        checkKeyLength(keyLength);
        Cipher cipher = Cipher.getInstance(ALGORITHM_NAME);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        byte[] input = new byte[length];
        System.arraycopy(cipherText, offset, input, 0, length);
        input = Base64.getDecoder().decode(input);
        int start = 0;
        length = input.length;
        byte[] buffer;
        int i = 0;
        final int MAX = getMaxDecryptLength(keyLength);
        while (length > start) {
            if (length - start > MAX) {
                buffer = cipher.doFinal(input, start, MAX);
            } else {
                buffer = cipher.doFinal(input, start, length - start);
            }
            outputStream.write(buffer, 0, buffer.length);
            i++;
            start = i * MAX;
        }
        input = outputStream.toByteArray();
        outputStream.close();
        return input;
    }

    /**
     * sign for data
     *
     * @param data       target data
     * @param privateKey sign by private key
     * @return the sign string
     */
    public static String sign(String data, PrivateKey privateKey) throws Exception {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKey.getEncoded());
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
        PrivateKey key = keyFactory.generatePrivate(keySpec);
        Signature signature = Signature.getInstance(MD5_RSA);
        signature.initSign(key);
        signature.update(data.getBytes());
        return Base64.getEncoder().encodeToString(signature.sign());
    }

    /**
     * verify the signature
     *
     * @param srcData   the data
     * @param publicKey verify by public key
     * @param sign      the sign string
     * @return if verify success
     */
    public static boolean verify(String srcData, PublicKey publicKey, String sign) throws Exception {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKey.getEncoded());
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
        PublicKey key = keyFactory.generatePublic(keySpec);
        Signature signature = Signature.getInstance(MD5_RSA);
        signature.initVerify(key);
        signature.update(srcData.getBytes());
        return signature.verify(Base64.getDecoder().decode(sign));
    }

    /**
     * get key pair
     *
     * @param keyLength the length of key
     * @return key pair
     */
    public static KeyPair getKeyPair(KeyLength keyLength) throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM_NAME);
        keyPairGenerator.initialize(getKeyLength(keyLength));
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * get public key by key string
     *
     * @param publicKeyString public key string
     * @return the key
     */
    public static PublicKey getPublicKey(String publicKeyString) throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
        byte[] decodedKey = Base64.getDecoder().decode(publicKeyString);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(decodedKey);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * get private key by key string
     *
     * @param privateKeyString private key string
     * @return the key
     */
    public static PrivateKey getPrivateKey(String privateKeyString) throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM_NAME);
        byte[] decodedKey = Base64.getDecoder().decode(privateKeyString);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(decodedKey);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * get public key in string
     *
     * @param keyPair key pair
     * @return the string of public key encoded by Base64
     */
    public static String getPublicKeyString(KeyPair keyPair) {
        if (keyPair == null) {
            throw new RuntimeException("KeyPair is null");
        }
        return Base64.getEncoder().encodeToString(keyPair.getPublic().getEncoded());
    }

    /**
     * get private key in string
     *
     * @param keyPair key pair
     * @return the string of private key encoded by Base64
     */
    public static String getPrivateKeyString(KeyPair keyPair) {
        if (keyPair == null) {
            throw new RuntimeException("KeyPair is null");
        }
        return Base64.getEncoder().encodeToString(keyPair.getPrivate().getEncoded());
    }

    /**
     * get length in bits
     *
     * @param keyLength enum KeyLength
     * @return length in bits
     */
    public static int getKeyLength(KeyLength keyLength) {
        switch (keyLength) {
            case LENGTH_256:
                return 256;
            case LENGTH_512:
                return 512;
            case LENGTH_1024:
                return 1024;
            case LENGTH_2048:
                return 2048;
            default:
                checkKeyLength(keyLength);
                return 1024;
        }
    }

    /**
     * get max length of encryption
     *
     * @return the max length
     */
    public static int getMaxEncryptLength(KeyLength keyLength) {
        switch (keyLength) {
            case LENGTH_256:
                return 21;
            case LENGTH_512:
                return 63;
            case LENGTH_1024:
                return 117;
            case LENGTH_2048:
                return 245;
            default:
                checkKeyLength(keyLength);
                return 117;
        }
    }

    /**
     * get max length of decryption
     *
     * @return the max length
     */
    public static int getMaxDecryptLength(KeyLength keyLength) {
        switch (keyLength) {
            case LENGTH_256:
                return 32;
            case LENGTH_512:
                return 64;
            case LENGTH_1024:
                return 128;
            case LENGTH_2048:
                return 256;
            default:
                checkKeyLength(keyLength);
                return 128;
        }
    }

    private static void checkKeyLength(KeyLength keyLength) {
        if (keyLength == null) {
            throw new RuntimeException("keyLength shouldn't be null");
        }
    }
}